/*
 * Copyright 2003 Program of Computer Graphics, Cornell University
 *     580 Rhodes Hall
 *     Cornell University
 *     Ithaca NY 14853
 * Web: http://www.graphics.cornell.edu/
 * 
 * Not for commercial use. Do not redistribute without permission.
 */

#ifndef _H_KDTREE
#define _H_KDTREE

//handle circularity in header file inclusion
class Scene;    

#include "Ray.h"
#include "Geometry.h"

#define KDTREE_MAX_DEPTH 56

//this class does not exist, instead we use it to indicate KDNode pointers with split type 
//encoded in the low two bits
class KDNodeHandle; 

//An interior node in the KDTree.  
class KDNode {
public:
  float splitValue; //splitvalue
  //note: split value (x,y,z, or leaf) is encoded in the pointer (handle) to this node
  class KDNodeHandle *left,*right; //handles to left and right children

  //if num==0 return left child, if num==1 return right child
  class KDNodeHandle *getChild(int num) const { return (&this->left)[num]; }
};

//These functions convert between a KDNodeHandle pointer and the KDNode pointer and
//split type (x,y,z, or leaf) that are encoded within a KDNodeHandle pointer
inline int getNodeType(KDNodeHandle *ptr) { return static_cast<int>( KDST_MASK&(size_t)ptr ); }
inline KDNode* getNodePtr(KDNodeHandle *ptr) { return (KDNode*)(((size_t)ptr)&~KDST_MASK); }
inline KDNodeHandle* makeKDNodeHandle(KDNode *ptr, int type)
{ assert((type== (type&KDST_MASK) )&&(type!=KDST_LEAF)); return (KDNodeHandle*)(type | (size_t)ptr); }
inline KDNodeHandle* makeKDNodeLeafHandle(PrimitiveList *listPtr)
{ return (KDNodeHandle*)(KDST_LEAF | (size_t)listPtr); }

//KDNodeArray is just a way of packing KDNode in cacheline size groups
#define KDNODEGROUP_ALIGNMENT 128
#define KDNODEGROUP_COUNT (KDNODEGROUP_ALIGNMENT/sizeof(KDNode))
class KDNodeGroup {
public:
  KDNode node[KDNODEGROUP_COUNT];
};

//kd node before tree compilation.  We first construct the KD-Tree using these and then
//compile/compress/pack the data into KDNodes for a smaller tree with better cache packing
class KDFatNode {
  int depthType;
public:
  float splitValue;
  class KDFatNode *left,*right;
  float area;

  int getSplitType() { return depthType&KDST_MASK; }
  int getDepth() { return depthType>>8; }
  void set(const KDSplitCode &code, int depth, const AxisAlignedBox &bounds) {
    depthType = (depth<<8)|(code.getSplitType()&KDST_MASK);
    splitValue = code.getSplitValue();
    left = right = NULL;
    F32vec4 size = bounds.max - bounds.min;
    area = size[0]*size[1] + size[1]*size[2] + size[2]*size[0];
  }
};

class KDTree {
  AxisAlignedBox epsBounds; //slightly enlarged bounding box
  AxisAlignedBox bounds;
  KDNodeHandle *rootHandle;
  KDNode dummy;
  KDFatNode *preliminaryRootPtr;
  PoolAllocator nodePool;
  PoolAllocator primitiveListPool;
  PoolAllocator fatNodePool;
  class KDNodeGroupCombiner *combiner;

  void buildNode(KDFatNode **res, class KDBuilder *builder, int depth, const AxisAlignedBox &bounds, GeometrySet &geom);
  void compileNodesSimplistic(KDNodeHandle **handle, KDFatNode *fat);
  void compileNodesChunkedBreadthFirst(KDNodeHandle **handle, KDFatNode *fat);
  void compileNodesChunkedMostLikely(KDNodeHandle **rootHandle, KDFatNode *rootFat);
public:
  KDTree(KDBuilder *buildStrategy, Scene *scene);
  ~KDTree();
  //override new and delete to force proper alignement
  static void* operator new (size_t size) { return malloc_aligned(static_cast<int>(size),16); }
  static void operator delete (void *p) { free_aligned(p); }

  //intersect routine is located in KDTreeTraverse.cpp
  void intersect(const Ray &ray, IntersectionRecord *hitRec);
  //intersect multiple rays (must be between 1 and 4 rays)
  void intersect(int numRays, const Ray ray[4], IntersectionRecord hitRec[4]);

  //get axis aligned bounding box
  AxisAlignedBox getBounds() { return bounds; }
  //get slightly enlarged version of bounding box (guard against roundoff error)
  AxisAlignedBox getBiggerBounds() { return epsBounds; } 
  //print node, leaf, and depth statistics about the current tree
  void printTreeStatistics();
};

//base class for objects encoding particular strategies for constructing the KDTree
class KDBuilder {
public:
  int buildPredLeftTri,buildPredRightTri; //for debugging only, predicted triangle splits
  int buildPredLeftOther,buildPredRightOther; //for debugging only, predicted other splits
  KDBuilder() { 
    buildPredLeftTri = -1; buildPredRightTri = -1; 
    buildPredLeftOther = -1; buildPredRightOther = -1; 
  }
  virtual ~KDBuilder() {};
  
  //decide whether to split a node or not and if so using which cut plane
  //return splitCode for node
  virtual KDSplitCode split(int depth, const AxisAlignedBox &bounds, const GeometrySet &geom)=0;
};

//some build strategies

class KDBFlatList : public KDBuilder {
public:
  KDBFlatList() {}
  virtual KDSplitCode split(int depth, const AxisAlignedBox &bounds, const GeometrySet &geom);
};

class KDBNaiveMiddleSplit : public KDBuilder {
  int maxDepth;
public:
  KDBNaiveMiddleSplit(int _maxDepth) { maxDepth = _maxDepth; }
  virtual KDSplitCode split(int depth, const AxisAlignedBox &bounds, const GeometrySet &geom);
};

class KDBNaiveMedianSplit : public KDBuilder {
  int maxElements;
public:
  KDBNaiveMedianSplit(int maxElementsPerLeaf) { maxElements = maxElementsPerLeaf; }
  virtual KDSplitCode split(int depth, const AxisAlignedBox &bounds, const GeometrySet &geom);
};

class KDBGreedyCostSplit : public KDBuilder {
private:
  KDFixedSplitEstimator fixed;
  float relSplitTable[KDFIXEDSPLIT_NUM]; //table of index/KDFIXEDSPLIT_NUM
  KDAllSplitEstimator all;

  KDSplitCode splitFixed(const AxisAlignedBox &bounds, const GeometrySet &geom);
  KDSplitCode splitAll(const AxisAlignedBox &bounds, const GeometrySet &geom);
public:
  KDBGreedyCostSplit();
  //override new and delete to force proper alignement
  static void* operator new (size_t size) { return malloc_aligned(static_cast<int>(size),16); }
  static void operator delete (void *p) { free_aligned(p); }

  virtual KDSplitCode split(int depth, const AxisAlignedBox &bounds, const GeometrySet &geom);
};


#endif //_H_KDTREE

